import h5py
import numpy as np
# import uuid
import os

import torch


class OriginData(object):
    def __init__(self):
        super().__init__()
        self.save_file_name = "/home/Dyf/code/dataset/st/origin_data.h5"
        self.features = []
        self.labels = []
        # for i in range(0,3):

    #     feat = features[i,:,:]
    #     name = str(i)
    #     write_h5(feat, name)
    def write_features(self, features):
        if not os.path.exists(self.save_file_name):
            last_id = 0
            mode = 'w'  # 如果文件不存在，则创建文件
            with h5py.File(self.save_file_name, mode) as file:
                for feature in features:
                    id = "{:010d}".format(last_id)
                    label_id = "label_" + id
                    file.create_dataset(name=id, data=feature[0])
                    file.create_dataset(name=label_id, data=feature[1])
                    last_id += 1
        else:
            last_id = self.get_last_id()
            mode = 'a'  # 如果文件不存在，则创建文件
            for feature in features:
                with h5py.File(self.save_file_name, mode) as file:
                    print("新增信息",feature[0],feature[1])
                    id = "{:010d}".format(last_id)
                    label_id = "label_" + id
                    file.create_dataset(name=id, data=feature[0])
                    file.create_dataset(name=label_id, data=feature[1])
                    last_id += 1

    def get_last_id(self):
        with h5py.File(self.save_file_name, 'r') as hf:
            last_id = len(hf.keys())
            return last_id

    def load_data(self):
        with h5py.File(self.save_file_name, 'r') as hf:
            for key in hf.keys():
                if "label_" in key:
                    label = torch.from_numpy(hf[key][:])
                    self.labels.append(label)
                else:
                    feature = torch.from_numpy(hf[key][:])

                    self.features.append(feature)
        # print(self.labels)
        # print(self.features)

    def __len__(self):
        with h5py.File(self.save_file_name, 'r') as hf:
            return len(hf.keys())

    def __getitem__(self, idx):
        input = self.features[idx]
        label = self.labels[idx]
        return input, label


ODData = OriginData()

#
# feature1 = np.array([[1, 4, 5, 6], [4, 8, 9, 0]], dtype=np.float32)
# label1 = np.array([1])
# feature2 = np.array([[2, 4, 5, 6], [5, 8, 9, 0]], dtype=np.float32)
# label2 = np.array([0])
# feature3 = np.array([[3, 4, 5, 6], [6, 8, 9, 0]], dtype=np.float32)
# label3 = np.array([1])
# ODData.write_features([(feature1, label1), (feature2, label2), (feature3, label3)])

# # print(od.features,"1")
# # print(len(od))
# # print(od)
# ODData.load_data()
# for i in ODData:
#     print(i)
# print(ODData[1])
